import os
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = "sk-w0P1gNv42gN6suR8819858Ee297c4b7cBc97Ae3e7b8dA701"
os.environ["OPENAI_API_BASE"] = "https://ai-yyds.com/v1"

# 指定ChromaDB的持久化目录
persist_directory = './mychroma_db'

# 加载ChromaDB
chroma_vector_database = Chroma(persist_directory=persist_directory)

# 查询问题
query = "美猴王最后回归花果山没有"
query_embedding = OpenAIEmbeddings(model='text-embedding-ada-002').embed_query(query)

# 在ChromaDB中搜索相似的文档块
results = chroma_vector_database.similarity_search(query_embedding, k=5)

print(results)

# 提取和显示相关文档块
# for result in results:
#     print(result.page_content)
#     print("\n---\n")